# letterbox方法resize

import os
import cv2
import numpy as np
import glob


def letterbox(img, new_shape=(640, 640), color=(114, 114, 114)):
    # Compute resizing dimensions
    shape = img.shape[:2]  # current shape [height, width]
    ratio = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    new_unpad = int(round(shape[1] * ratio)), int(round(shape[0] * ratio))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    dw /= 2  # divide padding by 2
    dh /= 2

    if shape[::-1] != new_unpad:  # resize
        img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = round(dh - 0.1), round(dh + 0.1)
    left, right = round(dw - 0.1), round(dw + 0.1)
    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)

    return img, ratio, dw, dh


def resize_label(bbox, img_shape, new_shape=(640, 640), ratio=1, dw=0, dh=0):
    x, y, w, h = bbox
    new_x = (x * img_shape[1] * ratio + dw) / new_shape[1]
    new_y = (y * img_shape[0] * ratio + dh) / new_shape[0]
    new_w = w * img_shape[1] * ratio / new_shape[1]
    new_h = h * img_shape[0] * ratio / new_shape[0]
    return new_x, new_y, new_w, new_h

def process_images(images_folder, labels_folder, output_folder):
    # Get list of image and label files
    image_files = sorted(glob.glob(os.path.join(images_folder, '*.jpg')))+ sorted(glob.glob(os.path.join(images_folder, '*.bmp')))
    label_files = sorted(glob.glob(os.path.join(labels_folder, '*.txt')))

    for image_file, label_file in zip(image_files, label_files):
        # Read image
        img = cv2.imread(image_file)
        img_shape = img.shape[:2]

        # Read label
        with open(label_file, 'r') as f:
            labels = f.readlines()
            labels = [list(map(float, label.strip().split())) for label in labels]

        # Resize image
        resized_img, ratio, dw, dh = letterbox(img)

        # Resize labels
        resized_labels = [resize_label(label[1:], img_shape, ratio=ratio, dw=dw, dh=dh) for label in labels]

        # Save resized image
        output_image_file = os.path.join(output_folder, 'images', os.path.basename(image_file))
        cv2.imwrite(output_image_file, resized_img)

        # Save resized labels
        output_label_file = os.path.join(output_folder, 'labels', os.path.basename(label_file))
        with open(output_label_file, 'w') as f:
            for label, resized_label in zip(labels, resized_labels):
                formatted_label = ' '.join(['{:.6f}'.format(coord) for coord in resized_label])
                f.write(f"{int(label[0])} {formatted_label}\n")


if __name__ == '__main__':
    # Set your folder paths here
    images_folder = 'F:/originData_train3/images/train'
    labels_folder = 'F:/originData_train3/labels/train'
    output_folder = 'F:/Resizedata/train'

    # Create output folders if they do not exist
    os.makedirs(os.path.join(output_folder, 'images'), exist_ok=True)
    os.makedirs(os.path.join(output_folder, 'labels'), exist_ok=True)

    process_images(images_folder, labels_folder, output_folder)

